In [1]:
    
from openmdao.main.api import Component, Assembly, Slot, Case
from openmdao.main.datatypes.api import Float, Int
from openmdao.util.testutil import assert_rel_error, assert_raises
from openmdao.lib.casehandlers.api import DBCaseRecorder, ListCaseRecorder, ListCaseIterator
from openmdao.lib.drivers.api import CaseIteratorDriver
from openmdao.main.api import Driver
from openmdao.lib.casehandlers.dbcase import list_db_vars, case_db_to_dict
from openmdao.lib.drivers.api import COBYLAdriver, CONMINdriver, \
        NEWSUMTdriver, SLSQPdriver, Genetic
    
import os.path as path
import pandas as pd
%matplotlib inline
from matplotlib.pylab import *
    
Some convenient function to manipulate openmdao classes
In [2]:
    
def call_func(self, outputs=None, **kwargs):
    """
    Conveniently call the assemblies/component in a functional way.
        output1, output2, ... = assembly(['output1', 'output2', ...], input1=val1, input2=val2...)
    or:
        output1 = assembly('output1', input1=val1, input2=val2...)
    or:
        assembly = assembly(input1=val1, input2=val2...)
    The function will do:
        self.input1 = val1
        self.input2 = val2
        ....
        self.execute()
        return [self.output1, self.output2,...]
    """
    for k, v in kwargs.iteritems():
        if hasattr(self, k):
            setattr(self, k, v)
            #print k + '=' + str(v)
    self.run()
    # Prepare outputs
    if not outputs:
        # Return the object executed that can then be used directly
        return self
    if isinstance(outputs, str):
        outputs_val = getattr(self, outputs)
    elif isinstance(outputs, list):
        outputs_val = []
        for o in outputs:
            outputs_val.append(getattr(self, o))
    return outputs_val
Component.__call__ = call_func
Assembly.__call__ = call_func
## Transform an openmdao recorder into a pandas dataframe
rec2df = lambda rec: pd.DataFrame([[c[k] for k in rec.cases[0].keys()] for c in rec.cases], columns=rec.cases[0].keys())
    
The Component: implement the Rosenbrock function: $$f=100 (x2 - x1^2)^2 + (1 - x1)^2$$
The Component has two inputs:
And one output:
In [4]:
    
class Rosenbrock(Component):
    """ Standard two-dimensional Rosenbrock function. """
    x1 = Float(iotype='in')
    x2 = Float(iotype='in')
    f  = Float(iotype='out')
    def execute(self):
        """ Just evaluate the function. """
        x1 = self.x1
        x2 = self.x2
        self.f = 100 * (x2 - x1**2)**2 + (1 - x1)**2
    
In [5]:
    
## Plotting the contour map
def contour_plot(Rosenbrock):
    rose = Rosenbrock()
    XS, YS = meshgrid(linspace(-2, 2, 20), linspace(-2,2, 20));
    ZS = array([rose(x1=x, x2=y).f for x,y in zip(XS.flatten(),YS.flatten())]).reshape(XS.shape);
    contourf(XS, YS, ZS, 50);
    colorbar()
contour_plot(Rosenbrock)
    
    
Looking at the Assembly structure:
In [52]:
    
class Optimization(Assembly):
    
    def configure(self):
        """ Configure driver and its workflow. """
        super(Assembly, self).configure()
        ## Add some components
        self.add('rosenbrock', Rosenbrock())
        ## Add an optimizer: COBYLAdriver, CONMINdriver, NEWSUMTdriver, SLSQPdriver, Genetic
        self.add('driver', Genetic())         
        self.driver.workflow.add('rosenbrock')
        
        ## The parameters of the optimization
        self.driver.add_parameter('rosenbrock.x1', low=-2, high=2, start=-1.0)
        self.driver.add_parameter('rosenbrock.x2', low=-2, high=2, start=-2.0)
        self.driver.add_objective('rosenbrock.f')
        
        ## This recorder will create a list of cases
        self.driver.recorders = [ListCaseRecorder()]
        ## These variables will be recorded
        self.driver.printvars = ['rosenbrock.x1','rosenbrock.x2','rosenbrock.f',]
        ## Some optimizer options
        self.driver.itmax = 300
        self.driver.fdch = 0.00001
        self.driver.fdchm = 0.000001
        self.driver.ctlmin = 0.001
        self.driver.delfun = 0.0001
    
In [53]:
    
%%timeit
opti = Optimization()
# optional here, replace optimizer: COBYLAdriver, CONMINdriver, NEWSUMTdriver, SLSQPdriver, Genetic
# opti.replace('driver', COBYLAdriver())
opti.run()
    
    
In [54]:
    
contour_plot(Rosenbrock)
## Transform the openmdao recorder into a pandas dataframe
df = rec2df(opti.driver.recorders[0])
df.plot(x='rosenbrock.x1', y='rosenbrock.x2', marker='x', color='w')
df[df['rosenbrock.f']==df['rosenbrock.f'].min()].plot(x='rosenbrock.x1', y='rosenbrock.x2', marker='o', color='r')
xlabel('x1'); ylabel('x2')
savefig('rosenbrock_optimization.pdf')
    
    
In [55]:
    
df.describe()
    
    Out[55]:
Recreating Rosenbrock as an Assembly of two Components:
In [72]:
    
class A(Component):
    x1 = Float(iotype='in')
    x2 = Float(iotype='in')
    
    x3 = Float(iotype='out')
    
    def execute(self):
        self.x3 = 100 * (self.x2 - self.x1**2)**2 
        
class B(Component):
    x1 = Float(iotype='in')
    x3 = Float(iotype='in')
    
    f = Float(iotype='out')
    
    def execute(self):
        self.f = self.x3 + (1 - self.x1)**2
    
class C(Assembly):
    x1 = Float(iotype='in')
    x2 = Float(iotype='in')
    f = Float(iotype='out')
    
    def configure(self):
        self.add('a', A())
        self.add('b', B())
        self.add('driver', Driver())
        
        ## The workflow is sequential: first `a` then `b`
        self.driver.workflow.add(['a', 'b'])
        ## Connecting the variables between the components and the assembly
        self.connect('a.x3', 'b.x3')
        self.connect('x1',['a.x1', 'b.x1'])
        self.connect('x2','a.x2')
        self.connect('b.f','f')        
contour_plot(C())
    
    
Now that we have recreated Rosenbrock, we can use it in our previous optimization assembly
In [73]:
    
opti = Optimization()
opti.configure()
## We replace the Rosenbrock instance by a C instance
opti.replace('rosenbrock', C())
opti.replace('driver', COBYLAdriver())
opti.execute()
    
    
In [3]:
    
contour_plot(Rosenbrock)
## Transform the openmdao recorder into a pandas dataframe
df = rec2df(opti.driver.recorders[0])
df.plot(x='rosenbrock.x1', y='rosenbrock.x2', marker='x', color='w')
df[df['rosenbrock.f']==df['rosenbrock.f'].min()].plot(x='rosenbrock.x1', y='rosenbrock.x2', marker='o', color='r')
xlabel('x1'); ylabel('x2')
savefig('rosenbrock_optimization.pdf')
    
    
In [ ]: